#!/usr/local/bin/python2.5

from math import pi
from time import localtime
from shelve import open as shopen
from pylab import *
from numpy import array, max, zeros, mean
from paUtils import *

colors = ['b','g', 'r', 'c', 'm', 'y', 'k', 'w']

class Plotter:
    def __init__(self, display, avg_spect):
        self.name = "Plotter"
        self.display = display
        self.avg_spect = avg_spect
        self.pN = 0
        self.cycle = 0
        self.cor = [0,0,0,0,0,0,0]
        self.pha = [0,0,0,0,0,0,0]
        self.mag = [0,0,0,0,0,0,0]
        self.fit = [0,0,0,0,0,0,0]
        self.fit_params = [0,0,0,0,0,0,0]
        self.pha_hist = {}
        self.del_hist = {}
        for i in range(7):
            self.pha_hist[i] = [0,]
            self.del_hist[i] = [0,]

    def loadData(self, data):
        if self.cycle<self.avg_spect:
            self.cycle += 1
        else:
            self.cycle = 0
            self.cor = [0,0,0,0,0,0,0]
        if len(data)>4:
            N, cor, pha, mag, fit, fit_params = data
            self.N = N
            for b in range(len(cor)):
                bias = mean(cor[b])
                if abs(bias)>2**29:
                    print "Overflow on %d? Bias is %d" %(b,bias)
                self.cor[b] += cor[b] - bias
            #self.pha = pha
            #self.mag = mag
            #self.fit = fit
            #self.fit_params = fit_params
            #self.addPhasePoint(self.fit_params)
        else:
            N, cor, pha, mag = data
            self.N = N
            self.cor = cor
            self.pha = pha
            self.mag = mag

    def addPhasePoint(self, fit_params):
        for i in range(len(fit_params)):
            self.del_hist[i].append(fit_params[i][0])
            self.pha_hist[i].append(fit_params[i][1])

    def plot(self, baselines):
        if self.cycle<>self.avg_spect:
            print "...averaging..."
            return
        print "...plotting..."
        for b in range(len(self.cor)):
            c = corrReorder(self.cor[b])
            self.pha[b] = getPhase(c)[2]
            self.mag[b] = getMagnitude(c)
            self.fit_params[b], self.fit[b] = getPhaseFit(dewrap(self.pha[b]))
        self.addPhasePoint(self.fit_params)
        if self.display=='baselines':
            self.plotBaselines(figure(1), baselines)
        if self.display=='correlations':
            self.plotCorrelations(baselines)
        if self.display=='phases':
            self.plotPhases(baselines)
        if self.display=='magnitudes':
            self.plotMagnitudes(baselines)
        

    def plotCorrelations(self, baselines=range(7)):
        if self.pN==0:
            fig = figure()
            self.total_splts = len(baselines)
            self.splts = {}
            self.lines = {}
            spnum = 0
            for i in baselines:
                spnum += 1
                self.splts[i] = fig.add_subplot(self.total_splts,
                                                1,
                                                spnum)
                self.lines[i] = self.splts[i].plot(range(-16,16),
                                              self.cor[i],
                                              colors[i]+"-",
                                              linewidth=2)[0]
        for i,line in self.lines.iteritems():
            line.set_ydata(self.cor[i])
            self.splts[i].set_ylim(min(self.cor[i]),max(self.cor[i]))
            self.splts[i].set_xlim(-16,16)
        draw()
        self.pN += 1

    def plotPhases(self, baselines=range(7)):
        if self.pN==0:
            fig = figure()
            self.total_splts = len(baselines)
            self.splts = {}
            self.lines = {}
            spnum = 0
            for i in baselines:
                spnum += 1
                self.splts[i] = fig.add_subplot(self.total_splts,
                                                1,
                                                spnum)
                self.lines[i] = self.splts[i].plot(range(0,16),
                                              self.pha[i],
                                              colors[i]+"h",
                                              markersize=10)[0]
        for i,line in self.lines.iteritems():
            line.set_ydata(self.pha[i])
            self.splts[i].set_ylim(-pi,pi)
            self.splts[i].set_xlim(0,16)
        draw()
        self.pN += 1

    def plotMagnitudes(self, baselines=range(7)):
        if self.pN==0:
            fig = figure()
            self.total_splts = len(baselines)
            self.splts = {}
            self.lines = {}
            spnum = 0
            for i in baselines:
                spnum += 1
                self.splts[i] = fig.add_subplot(self.total_splts,
                                                1,
                                                spnum)
                self.lines[i] = self.splts[i].plot(range(0,16),
                                              self.mag[i],
                                              colors[i]+"-",
                                              linewidth=4)[0]
        for i,line in self.lines.iteritems():
            line.set_ydata(self.mag[i])
            self.splts[i].set_ylim(min(self.mag[i]),max(self.mag[i]))
            self.splts[i].set_xlim(0,16)
        draw()
        self.pN += 1

    def plotData(self, fig, baselines=range(7)):
        if self.pN==0:
            # fig = figure()
            self.total_splts = len(baselines)
            self.splts = {}
            self.lines = {}
            spnum = 0
            for i in baselines:
                spnum += 1
                self.splts[i] = [ fig.add_subplot(self.total_splts,
                                                  2,
                                                  spnum),
                                  fig.add_subplot(self.total_splts,
                                                  2,
                                                  spnum+1) ]
                spnum += 1
                self.lines[i] = [ self.splts[i][0].plot(range(-16,16),
                                                        self.cor[i],
                                                        colors[i]+"-",
                                                        linewidth=2)[0],
                                  self.splts[i][1].plot(range(1,16),
                                                        self.pha[i][1:],
                                                        colors[i]+"h",
                                                        markersize=10)[0],
                                  self.splts[i][1].plot(range(1,16),
                                                        self.mag[i][1:],
                                                        colors[i]+'-',
                                                        linewidth=4)[0] ]
        maxall = 0
        for i in baselines:
            if max(self.mag[i][1:])>maxall:
                maxall = max(self.mag[i][1:])
        if maxall==0: maxall=1
        for i,line in self.lines.iteritems():
            line[0].set_ydata(self.cor[i])
            line[1].set_ydata(self.pha[i][1:])
            line[2].set_ydata(pi*array(self.mag[i][1:])/maxall)
            self.splts[i][0].set_ylim(min(self.cor[i]),max(self.cor[i]))
            self.splts[i][0].set_xlim(-16,16)
            self.splts[i][1].set_ylim(-pi*1.25,pi*1.25)
            self.splts[i][1].set_xlim(0,16)
        draw()
        self.pN += 1

    def plotBaselines(self, fig, baselines=range(7), n=-1):
        corticks = [-16, -12, -8, -4, 0, 4, 8, 12, 16]
        corticklabel = ['-16', '', '-8', '', '0', '', '8', '', '16']
        phaticks = [0, 2, 4, 6, 8, 10, 12, 14, 16]
        phaticklabel = ['0', '', '128', '', '256', '', '385', '', '512']
        phaticks.reverse()
        histicks = [0, 5, 10, 15, 20, 25, 30]
        histicklabel = ['-16', '', '', '-8', '', '', 'Last Scan']
        if n==-1: n=self.pN
        if (n-30)<0:
            start_hist = 0
            stop_hist = 30
        else:
            start_hist = n-30
            stop_hist = n
        if self.pN==0:
            #fig = figure()
            self.total_splts = len(baselines)
            self.splts = {}
            self.lines = {}
            spnum = 0
            for i in baselines:
                spnum += 1
                self.splts[i] = [fig.add_subplot(
                    self.total_splts,
                    3,
                    spnum,
                    xticks=corticks,
                    xticklabels=[]\
                    or corticklabel*(i==6),
                    yticklabels=[],
                    xlabel="Correlator Lags",
                    ylabel="6-%s"%((i+1)*(i<5) or 7*(i==5) or 8*(i==6))),
                                  fig.add_subplot(\
                    self.total_splts,
                    3,
                    spnum+1,
                    xticks=phaticks,
                    xticklabels=[]\
                    or phaticklabel*(i==6),
                    yticklabels=[],
                    xlabel="Frequency (MHz)",
                    ylabel="" or "Phase (dots)"*(i==5) or \
                                                 "Amplitudes (lines)"*(i==1)),
                                  fig.add_subplot(\
                    self.total_splts,
                    3,
                    spnum+2,
                    xticks=histicks,
                    xticklabels=[]\
                    or histicklabel*(i==6),
                    yticklabels=[],
                    xlabel="Fitting History (mins)",
                    ylabel="" or "Phase fit (dots)"*(i==5) or \
                                                 "Delay fit (lines)"*(i==1)) ]
                spnum += 2
                self.lines[i] = [ self.splts[i][0].plot(range(-16,16),
                                                        self.cor[i],
                                                        colors[i]+"-",
                                                        linewidth=2)[0],
                                  self.splts[i][1].plot(range(1,16),
                                                        self.pha[i][1:],
                                                        colors[i]+"h",
                                                        markersize=10)[0],
                                  self.splts[i][1].plot(range(1,16),
                                                        self.mag[i][1:],
                                                        colors[i]+'-',
                                                        linewidth=4)[0],
                                  self.splts[i][1].plot(range(0,16),
                                                        self.fit[i],
                                                        colors[i]+'-',
                                                        linewidth=1)[0],
                                  self.splts[i][2].plot(range(0,len(self.pha_hist[i][start_hist:stop_hist])),
                                                        self.pha_hist[i][start_hist:stop_hist],
                                                        colors[i]+'h',
                                                        markersize=3)[0],
                                  self.splts[i][2].plot(range(0,len(self.del_hist[i][start_hist:stop_hist])),
                                                        self.del_hist[i][start_hist:stop_hist],
                                                        colors[i]+'-',
                                                        linewidth=1)[0] ]
        maxall = 0.0
        #maxdel = {}
        #mindel = {}
        #normdel = {}
        for i in baselines:
            #self.del_hist[i].append(self.fit_params[i][0])
            #self.pha_hist[i].append(self.fit_params[i][1])
            #maxdel[i] = max(self.del_hist[i][-30:])
            #mindel[i] = min(self.del_hist[i][-30:])
            #normdel[i] = max([abs(maxdel[i]),abs(mindel[i])])
            #if normdel[i]==0: normdel[i]=1
            if max(self.mag[i][1:])>maxall:
                maxall = max(self.mag[i][1:])
        if maxall==0: maxall=1
        for i,line in self.lines.iteritems():
            line[0].set_ydata(self.cor[i])
            line[1].set_ydata(self.pha[i][1:])
            line[2].set_ydata(pi*array(self.mag[i][1:])/maxall)
            line[3].set_ydata(self.fit[i])
            line[4].set_xdata(range(len(self.pha_hist[i][start_hist:stop_hist])))
            line[4].set_ydata(self.pha_hist[i][start_hist:stop_hist])
            line[5].set_xdata(range(len(self.del_hist[i][start_hist:stop_hist])))
            line[5].set_ydata(pi*array(self.del_hist[i][start_hist:stop_hist]))
            self.splts[i][0].set_ylim(min(self.cor[i]),max(self.cor[i]))
            self.splts[i][0].set_xlim(-16,16)
            self.splts[i][1].set_ylim(-pi*1.25,pi*1.25)
            self.splts[i][1].set_xlim(0,16)
            self.splts[i][2].set_ylim(-pi*1.25,pi*1.25)
            self.splts[i][2].set_xlim(0,len(self.pha_hist[i][start_hist:stop_hist]))
        draw()
        self.pN += 1

def formatTime():
    (year, month, day,
    hours, minutes, seconds,
    weekday, jd, dst) = localtime()
    date = '/'.join((str(day),str(month),str(year)))
    time = ':'.join((str(hours),str(minutes),str(seconds)))
    return date+' at '+time

def unformatTime(date_str):
    date_str, time_str = date_str.split(' at ')
    day_str, mon_str, year_str = date_str.split('/')
    hr_str, min_str, sec_str = time_str.split(':')
    return int(year_str),int(mon_str),int(day_str),\
        int(hr_str),int(min_str),int(sec_str)

def cmpTime(x,y):
    if x==y:
        return 0
    for i in range(len(x)):
        if x[i]>y[i]:
            return 1
        elif y[i]>x[i]:
            return -1

if __name__ == "__main__":
    from sys import argv
    from optparse import OptionParser
    from time import time, sleep
    from rpyc import connect
    usage = "usage: %prog [options] arg"
    parser = OptionParser(usage)
    parser.add_option("-r",
                      "--remote",
                      action = "store_true",
                      dest = "remote")
    parser.add_option("-c",
                      "--host",
                      action = "store",
                      type = "string",
                      dest = "host",
                      default = "localhost")
    parser.add_option("-p",
                      "--port",
                      action = "store",
                      type = "int",
                      dest = "port",
                      default = 18861)
    parser.add_option("-b",
                      "--baselines",
                      action = "store",
                      type = "string",
                      dest = "baselines",
                      default = "0,1,2,3")
    parser.add_option("-i",
                      "--integration-time",
                      action = "store",
                      type = "int",
                      dest = "int_time",
                      default = 8)
    parser.add_option("-a",
                      "--average-spectra",
                      action = "store",
                      type = "int",
                      dest = "avg_spect",
                      default = 4)
    parser.add_option("-l",
                      "--log-filename",
                      action = "store",
                      type = "string",
                      dest = "logfilename")
    parser.add_option("-d",
                      "--display",
                      action = "store",
                      type = "string",
                      dest = "display",
                      default = "None")
    (options, args) = parser.parse_args()
    host = options.host
    port = options.port
    baselines = options.baselines.split(",")
    for i in range(len(baselines)):
        baselines[i] = int(baselines[i])
    int_time = options.int_time
    avg_spect = options.avg_spect
    display = options.display
    pa = connect(host,port)
    pap = Plotter(display, avg_spect)
    if options.logfilename:
        logfile = shopen(options.logfilename)
        logfile['start'] = formatTime()
        print "Created logfile at "+options.logfilename
    ion()
    print "Ctrl-C to stop..."
    while True:
        if options.remote:
            data = pa.root.padata()
        else:
            data = pa.root.startCorrelation(int_time)
        tstart = time()
        if display!='None':
            pap.loadData(data)
            print pap.cycle
            pap.plot(baselines)
        if options.logfilename:
            logfile[formatTime()] = data[0:4]
        if (time()-tstart)>int_time: print "Plotting is taking too long!"
        while (time()-tstart)<int_time:
            tnow = time()
        if not options.remote: pa.root.retrieveCorrelation()
        print "Period: "+str(time()-tstart)
    if options.logfilename:
        print "Data stored to "+options.logfilename+":"
        for k in logfile.keys():
            print k
        logfile.close()
    print "Stopped."
